""" SynergyNet infer script. """
import argparse
import mindspore
import cv2
import numpy as np
import sys

from mindspore import context, Tensor, load_checkpoint, load_param_into_net
import mindspore.dataset.vision as vision
import mindspore.ops as ops

from mind3d.utils.synergynet_builder import build_model
from mind3d.utils.synergynet_util import Crop
from mind3d.utils.load_yaml import load_yaml

from mind3d.utils.synergynet_util import ParamsPack

param_pack = ParamsPack()


def img_loader(path):
    return cv2.imread(path, cv2.IMREAD_COLOR)


def infer(opt):
    """Main funtion for the infering process"""
    print("infer...")
    context.set_context(max_call_depth=100000, device_id=opt['device_id'], mode=context.GRAPH_MODE,
                        device_target=opt['device_target'])

    print('loading model...')
    model = build_model(img_size=opt['infer']['img_size'], mode="test")
    checkpoint = load_checkpoint(opt['infer']['checkpoint_fp'])
    load_param_into_net(model, checkpoint)

    img = np.array(img_loader(opt['infer']['path'])).astype(np.float32)

    transpose = vision.HWC2CHW()
    crop = Crop(5, mode='test')
    mean_channel = [127.5]
    std_channel = [128]
    normalize_op = vision.Normalize(mean=mean_channel, std=std_channel)
    trans = mindspore.dataset.transforms.Compose([transpose, crop, normalize_op])

    img = trans(img)
    img = Tensor(img, dtype=mindspore.float32)
    expand_dims = ops.ExpandDims()
    img = expand_dims(img, 0)

    print(img.shape)
    output = model(img)
    param_prediction = output.asnumpy()
    print(param_prediction)
    print('infer completed...')


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='SynergyNet train')
    parser.add_argument('-opt', type=str, default='mind3d/configs/synergy_net/synergynet.yaml',
                        help='Path to option YAML file.')
    args = parser.parse_known_args()[0]
    opt = load_yaml(args.opt)

    infer(opt)
